import numpy as np
from pysr import PySRRegressor
from bishe_benchmark import get_all_nguyen_equations, generate_nguyen_data
from bishe_situations.only_pysr import PysrTrainer
from bishe_situations.only_gplearn import GpLearnTrainer
from bishe_situations.neural_gplearn import NeuralGplearnTrainer
from bishe_situations.neural_pysr import NeuralPysrTrainer
import torch
import time
from typing import Dict
import json
import os, pandas as pd

def use_NeuralGPlearn(X, y):
    model = NeuralGplearnTrainer(
        population_size=50, 
        generations=16,
        niterations=50,
        device='cuda:0'
    )
    model.fit(X, y)
    return model.best_function(), model.mse, model.train_time

def use_GPlearn(X, y):
    # 初始化PySR模型
    model = GpLearnTrainer(
        population_size=50, 
        generations=16,
    )
    # 训练模型
    model.fit(X, y)
    return model.best_function(), model.mse, model.train_time

def use_NeuralPySR(X, y):
    model = NeuralPysrTrainer(
        population_size=50, 
        generations=16,
        niterations=50,
        device='cuda:1'
    )
    model.fit(X, y)
    return model.best_function(), model.mse, model.train_time

def use_PySR(X, y):
    # 初始化PySR模型
    model = PysrTrainer(
        population_size=50, 
        generations=16,
    )
    # 训练模型
    model.fit(X, y)
    return model.best_function(), model.mse, model.train_time


def run_benchmark(
    save_dir: str = "benchmark_results"
) -> Dict:
    """运行基准测试"""

    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)
    
    # 获取所有Nguyen方程
    equations = get_all_nguyen_equations()
    results = {}
    
    t = time.strftime("%Y%m%d_%H%M%S", time.localtime())
    root_file_prefix = "arbitary/" + f"{t}/"
    for func, name in equations:
        print(f"\n开始测试方程: {name}")
        root_dir = root_file_prefix + f"{name}/"
        os.makedirs(root_dir, exist_ok=True)
        
        # 生成数据
        X, y = generate_nguyen_data(func, n_samples=1000, noise=0.1)
        X = X.reshape(-1, 1)
        y = y.reshape(-1, 1)
        
        
        with open(root_dir + "result.csv", "w", encoding="utf-8") as f:
            # 优化超参数
            f.write(f"name,mse,train_time,func\n")
            # start NeuralPySR
            print("start NeuralPySR")
            _func, mse, train_time, mse_list = use_NeuralPySR(X, y)
            # print(func, mse, train_time, mse_list)
            f.write(f"{name},{mse},{train_time},{_func}\n")
            df = pd.DataFrame(mse_list)
            df.to_csv(root_dir + "NeuralPySR_mse_list.csv", index=False)
            # 使用PySR
            print("start PySR")
            _func, mse, train_time = use_PySR(X, y)
            # print(func, mse, train_time)
            f.write(f"{name},{mse},{train_time},{_func}\n")
            # 使用NeuralGPlearn
            print("start NeuralGPlearn")
            _func, mse, train_time, mse_list = use_NeuralGPlearn(X, y)
            # print(func, mse, train_time, mse_list)
            f.write(f"{name},{mse},{train_time},{_func}\n")
            df = pd.DataFrame(mse_list)
            df.to_csv(root_dir + "NeuralGPlearn_mse_list.csv", index=False)
            # 使用GPlearn
            print("start GPlearn")
            _func, mse, train_time = use_GPlearn(X, y)
            # print(func, mse, train_time)
            f.write(f"{name},{mse},{train_time},{_func}\n")
    
    return results


if __name__ == "__main__":
    """主函数"""
    # 设置随机种子
    np.random.seed(42)
    torch.manual_seed(42)
    
    # 运行基准测试
    results = run_benchmark(save_dir="benchmark_results")